"""
LangGraph核心状态和配置管理
"""
from typing import TypedDict, Annotated, List, Dict, Any, Optional
from datetime import datetime
import uuid


class GraphState(TypedDict):
    """LangGraph状态定义"""
    # 对话相关
    conversation_id: Optional[int]
    user_id: Optional[int]
    
    # 消息历史
    messages: List[Dict[str, str]]
    
    # LLM配置
    provider: str
    model_name: str
    temperature: float
    max_tokens: int
    
    # 工作流控制
    current_step: str
    next_step: str
    
    # 工具调用
    tools_used: List[str]
    function_calls: List[Dict[str, Any]]
    
    # 错误处理
    error_occurred: bool
    error_message: Optional[str]
    retry_count: int
    
    # 性能指标
    start_time: datetime
    end_time: Optional[datetime]
    total_tokens: int
    estimated_cost: float
    
    # 上下文数据
    context_data: Dict[str, Any]
    
    # 工作流元数据
    workflow_id: str
    workflow_type: str
    
    # 流式响应控制
    is_streaming: bool
    stream_buffer: List[str]


class WorkflowConfig:
    """工作流配置"""
    
    def __init__(
        self,
        workflow_type: str = "simple_chat",
        max_retries: int = 3,
        timeout_seconds: int = 300,
        enable_tools: bool = True,
        enable_memory: bool = True,
        enable_streaming: bool = True
    ):
        self.workflow_type = workflow_type
        self.max_retries = max_retries
        self.timeout_seconds = timeout_seconds
        self.enable_tools = enable_tools
        self.enable_memory = enable_memory
        self.enable_streaming = enable_streaming
        
        # 工作流特定配置
        self.workflow_configs = {
            "simple_chat": {
                "steps": ["validate_input", "generate_response"],
                "tools": []
            },
            "tool_enhanced": {
                "steps": ["validate_input", "tool_selection", "tool_execution", "generate_response"],
                "tools": ["search", "calculator", "code_executor"]
            },
            "reasoning": {
                "steps": ["validate_input", "reasoning_chain", "generate_response"],
                "tools": []
            }
        }
    
    def get_workflow_steps(self) -> List[str]:
        """获取工作流步骤"""
        return self.workflow_configs.get(self.workflow_type, {}).get("steps", [])
    
    def get_available_tools(self) -> List[str]:
        """获取可用工具"""
        return self.workflow_configs.get(self.workflow_type, {}).get("tools", [])


class StateManager:
    """状态管理器"""
    
    def __init__(self):
        self.active_states = {}
    
    def create_state(self, initial_data: Dict[str, Any]) -> GraphState:
        """创建新状态"""
        state_id = str(uuid.uuid4())
        
        state = GraphState(
            conversation_id=initial_data.get("conversation_id"),
            user_id=initial_data.get("user_id"),
            messages=initial_data.get("messages", []),
            provider=initial_data.get("provider", "openai"),
            model_name=initial_data.get("model_name", "gpt-3.5-turbo"),
            temperature=initial_data.get("temperature", 0.7),
            max_tokens=initial_data.get("max_tokens", 4000),
            current_step="start",
            next_step="validate_input",
            tools_used=[],
            function_calls=[],
            error_occurred=False,
            error_message=None,
            retry_count=0,
            start_time=datetime.now(),
            end_time=None,
            total_tokens=0,
            estimated_cost=0.0,
            context_data=initial_data.get("context_data", {}),
            workflow_id=state_id,
            workflow_type=initial_data.get("workflow_type", "simple_chat"),
            is_streaming=initial_data.get("is_streaming", False),
            stream_buffer=[]
        )
        
        self.active_states[state_id] = state
        return state
    
    def get_state(self, state_id: str) -> Optional[GraphState]:
        """获取状态"""
        return self.active_states.get(state_id)
    
    def update_state(self, state_id: str, updates: Dict[str, Any]) -> bool:
        """更新状态"""
        if state_id not in self.active_states:
            return False
        
        state = self.active_states[state_id]
        for key, value in updates.items():
            if hasattr(state, key):
                setattr(state, key, value)
        
        return True
    
    def cleanup_state(self, state_id: str):
        """清理状态"""
        if state_id in self.active_states:
            del self.active_states[state_id]